#include "VTKWriter.h"
#include <fstream>
#include <iomanip>

void VTKWriter::writeHeader(std::ofstream& file, size_t numPoints, size_t numCells) {
    file << "# vtk DataFile Version 3.0\n";
    file << "Corner Grid Model Surfaces\n";
    file << "ASCII\n";
    file << "DATASET UNSTRUCTURED_GRID\n\n";
}

void VTKWriter::writePoints(std::ofstream& file, const std::vector<Point3D>& vertices) {
    file << "POINTS " << vertices.size() << " double\n";
    for (const auto& vertex : vertices) {
        file << std::fixed << std::setprecision(6)
             << vertex.x << " " << vertex.y << " " << vertex.z << "\n";
    }
    file << "\n";
}

void VTKWriter::writeCells(std::ofstream& file, 
                          const std::vector<Triangle>& strataTriangles,
                          const std::vector<Quad>& faultQuads,
                          const std::vector<Quad>& strataQuads) {
    // 计算所有单元的总数
    size_t totalCells = strataTriangles.size() + faultQuads.size() + strataQuads.size();
    size_t totalIndices = strataTriangles.size() * 4 + (faultQuads.size() + strataQuads.size()) * 5;
    
    // 写入单元连接性
    file << "CELLS " << totalCells << " " << totalIndices << "\n";
    
    // 写入三角形单元
    for (const auto& tri : strataTriangles) {
        file << "3 " << tri.v1 << " " << tri.v2 << " " << tri.v3 << "\n";
    }
    
    // 写入断层四边形单元
    for (const auto& quad : faultQuads) {
        file << "4 " << quad.v1 << " " << quad.v2 << " " << quad.v3 << " " << quad.v4 << "\n";
    }
    
    // 写入地层四边形单元
    for (const auto& quad : strataQuads) {
        file << "4 " << quad.v1 << " " << quad.v2 << " " << quad.v3 << " " << quad.v4 << "\n";
    }
    file << "\n";

    // 写入单元类型
    file << "CELL_TYPES " << totalCells << "\n";
    // VTK_TRIANGLE = 5, VTK_QUAD = 9
    for (size_t i = 0; i < strataTriangles.size(); ++i) {
        file << "5\n";  // VTK_TRIANGLE
    }
    for (size_t i = 0; i < faultQuads.size() + strataQuads.size(); ++i) {
        file << "9\n";  // VTK_QUAD
    }
    file << "\n";
}

void VTKWriter::writePointData(std::ofstream& file, const std::vector<Point3D>& vertices) {
    file << "POINT_DATA " << vertices.size() << "\n";
    file << "SCALARS IsInterface int 1\n";
    file << "LOOKUP_TABLE default\n";
    for (const auto& vertex : vertices) {
        file << vertex.type << "\n";  // 断层节点type为1-10，地层节点type为11-20
    }
    file << "\n";
}

void VTKWriter::writeVTKFile(const std::string& filename, const SurfaceExtractor& extractor) {
    std::ofstream file(filename);
    if (!file) {
        throw std::runtime_error("无法创建VTK文件: " + filename);
    }

    const auto& vertices = extractor.getVertices();
    const auto& strataTriangles = extractor.getStrataTriangles();
    const auto& faultQuads = extractor.getFaultQuads();
    const auto& strataQuads = extractor.getStrataQuads();

    size_t numCells = strataTriangles.size() + faultQuads.size() + strataQuads.size();

    // 写入文件各个部分
    writeHeader(file, vertices.size(), numCells);
    writePoints(file, vertices);
    writeCells(file, strataTriangles, faultQuads, strataQuads);
    writePointData(file, vertices);

    file.close();
} 